/*
 * Variable length code text compressor
 * (c) Copyright 2010 ccr/TNSP <ccr@tnsp.org>
 */
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <ctype.h>

void printbits(FILE *f, int val, int n)
{
  int i;
  for (i = n - 1; i >= 0; i--) {
    fputc((val & (1 << i)) ? '1' : '0', f);
  }
}

typedef struct {
  FILE *f;
  int buf, bitcnt, bytecnt;
} putbits_ctx;


int putbits_init(putbits_ctx *ctx, FILE *f)
{
  if (f == NULL || ctx == NULL)
    return -1;
  ctx->f = f;
  ctx->buf = 0;
  ctx->bytecnt = 0;
  ctx->bitcnt = 8;
  return 0;
}

void putbits(putbits_ctx *ctx, int val, int n)
{
  int i;
  unsigned int mask = 1 << (n - 1);

  for (i = 0; i < n; i++) {
    ctx->buf <<= 1;

    if (val & mask)
      ctx->buf |= 1;

    mask >>= 1;
    ctx->bitcnt--;

    if (ctx->bitcnt == 0)
    {
      fputc((ctx->buf & 0xff), ctx->f);
      ctx->bitcnt = 8;
      ctx->bytecnt++;
    }
  }
}

int putbits_close(putbits_ctx *ctx)
{
  if (ctx == NULL) return -1;

  if (ctx->bitcnt != 8)
    putbits(ctx, 0, ctx->bitcnt);
  
  return 0;
}


typedef struct {
  int c, freq, code, len;
} freq_t;


int analyze_cmpfunc(const void *p1, const void *p2)
{
  const freq_t *s1 = p1, *s2 = p2;
  return s2->freq - s1->freq;
}

#define MAXTAB (256)
freq_t freqs[MAXTAB];

void analyze_buf(char *buf, int len)
{
  int k, n, maxcodes;

  /* Calculate frequencies */
  fprintf(stderr, "* Building frequency table %p, %d.\n", buf, len);
  memset(freqs, 0, sizeof(freqs));

  for (n = 0; n < len; n++) {
    int c = tolower(buf[n]);
    if (c == '\n' || c == '\r') c = ' ';
    buf[n] = c;
    if (c >= 0 && c < MAXTAB)
      freqs[c].freq++;
    else
      fprintf(stderr, "ERR %d!", c);
  }

  /* Count number of codes */
  for (maxcodes = n = 0; n < MAXTAB; n++) {
    freqs[n].c = n;
    if (freqs[n].freq > 0) maxcodes++;
  }
  fprintf(stderr, "* Total codes = %d\n", maxcodes);

  qsort(freqs, MAXTAB, sizeof(freq_t), analyze_cmpfunc);

  for (k = n = 0; n < MAXTAB; n++)
  if (freqs[n].freq > 0) {
    int c = -1, len;

    if (k == 0) 	{ c = 0; len = 2; }
    else
    if (k <= 4)		{ c = 4 + (k - 1); len = 4; }
    else
    if (k <= 12)	{ c = 16 + (k - 5); len = 5; }
    else
    if (k < maxcodes)	{ c = 48 + (k - 13); len = 6; }
    
    if (c < 0) {
      fprintf(stderr, "Internal error in codeforming: n=%d, k=%d\n", n, k);
      exit(5);
    }
    
    freqs[n].code = c;
    freqs[n].len = len;
    
    fprintf(stderr, "[%3d:%3d:%02x:'%c'] ", n, freqs[n].freq, freqs[n].c, freqs[n].c); printbits(stderr, c, len); fprintf(stderr, "\n");
    k++;
  }

}

#define BUFSIZE 8192

int main(int argc, char *argv[])
{
  char *infilename, *outfilename;
  FILE *f;
  putbits_ctx pb;
  char buf[BUFSIZE];
  int len, i, total;

  if (argc < 3) {
    fprintf(stderr, "Usage: %s <infile> <outfile>\n", argv[0]);
    return 1;
  }
  
  infilename = argv[1];
  outfilename = argv[2];

  fprintf(stderr, "* Reading input '%s'\n", infilename);
  if ((f = fopen(infilename, "rb")) == NULL) {
    fprintf(stderr, "ERROR: Could not open input file '%s'\n", infilename);
    return -1;
  }
  
  len = fread(buf, sizeof(char), BUFSIZE, f);
  fclose(f);

  analyze_buf(buf, len);

  /* Encoded data to specified output file
   */
  fprintf(stderr, "* Writing encoded output to '%s'\n", outfilename);
  if ((f = fopen(outfilename, "wb")) == NULL) {
    fprintf(stderr, "ERROR: Could not open output file '%s'\n", outfilename);
    return -1;
  }
  putbits_init(&pb, f);
  for (total = i = 0; i < len; i++) {
    int n, v = buf[i];
    for (n = 0; n < MAXTAB; n++) {
      if (freqs[n].c == v) {
        putbits(&pb, freqs[n].code, freqs[n].len);
        total += freqs[n].len;
      }
    }
  }
  putbits_close(&pb);
  fclose(f);

  /* Output codetable to stdout
   */
  printf("codetab: .byte ");
  fprintf(stderr, "#");
  for (i = 1; i < MAXTAB; i++)
  if (freqs[i].freq > 0) {
    int q = freqs[i].c;
    if (q >= 'a' && q <= 'z')
      q = q - 'a' + 1;
    printf("$%02x, ", q);
    fprintf(stderr, "%02x ", q);
  }
  fprintf(stderr, "\n");
  printf("\n");

  /* Display some statistics
   */
  fprintf(stderr, "* %d bytes (%d bits) in - %1.2f bytes (%d bits) out, %1.2f ratio.\n",
    len, len * 8,
    ((double) total) / 8.0f, total,
    ((double) total) / ((double) len * 8.0));
  
  return 0;
}
